import numpy as np
import cv2
import glob
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

# 设置棋盘格十字交叉的行数和列数
chessboard_size = (6, 4)
#写入示例图像的位置
exsample_img = '2024-03-10 1606236.png'
#图像上移的像素点个数，因为可能底部有全黑，对控制不利
img_shift_up = 13
#图像右移像素点个数，用来中心对称（说实话如果它不是0该考虑一下是不是摄像头安装位置有问题）
img_shift_right = 0
#alpha是个0~1之间的浮点数，其代表了无效像素点的保留量，1表示原图所有像素点全部保留，0表示仅保留有效像素点
#详见博客https://www.cnblogs.com/riddick/p/6711263.html
alpha = 0.5






# 准备棋盘格角点的坐标
objp = np.zeros((chessboard_size[0] * chessboard_size[1], 3), np.float32)
objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2)

# 存储棋盘格图像的角点坐标和三维空间中的坐标
objpoints = []  # 三维空间中的棋盘格角点坐标
imgpoints = []  # 图像平面的棋盘格角点坐标

# 加载所有的棋盘格图像
images = glob.glob('chess_board/*.png')

for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 寻找棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

    if ret:
        objpoints.append(objp)
        imgpoints.append(corners)

        # 可视化检测到的角点
        cv2.drawChessboardCorners(img, chessboard_size, corners, ret)
        cv2.imshow('Chessboard Corners', img)
        cv2.waitKey(50)  # 显示图像的时间，单位为毫秒

cv2.destroyAllWindows()

# 相机标定
ret, cameraMatrix, distCoeffs, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)

# 读取一张图像并去畸变校正
# img = cv2.imread('test_image.jpg')
h, w = img.shape[:2]

# 获取优化后的新相机矩阵
newCameraMatrix, roi = cv2.getOptimalNewCameraMatrix(cameraMatrix, distCoeffs, (w,h), alpha, (w,h),True)


newCameraMatrix[1,2] += img_shift_up 
newCameraMatrix[0,2] += img_shift_right 


img = cv2.imread(exsample_img)
# 使用undistort函数校正图像


dst = cv2.undistort(img, cameraMatrix, distCoeffs, None, newCameraMatrix)
# 将校正后的图像放大到原始图像的n倍大小
dst_resized = cv2.resize(dst, None, fx=1, fy=1, interpolation=cv2.INTER_LINEAR)
# 显示原始图像和校正后的图像

#在控制台输出变换矩阵
print("New Camera Matrix:")
print(newCameraMatrix)
#########################

# 创建两张子图
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))  # 1 行 2 列的布局

ax1.imshow(img)
ax1.set_title('Origin')
ax1.axis('off') 

ax2.imshow(dst_resized)
ax2.set_title('Undistorted')
ax2.axis('off') 

# 调整子图之间的间距
plt.tight_layout()
plt.show()






